﻿# -*- coding: utf-8 -*-
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from scipy.special import softmax
import random
import math
  
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows

class UnionFind:
    def __init__(self, elements):
        self.parent = {}   # u:v  u的父节点是v,通过v找到u的根节点
        for e in elements: 
            self.parent[e] = e
    
    def find(self, u):  # 查找u的根节点
        if self.parent[u] != u:
            self.parent[u] = self.find(self.parent[u])  # 路径压缩
        return self.parent[u]
    
    def union(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)
        if root_u != root_v:
            self.parent[root_u] = root_v #  将u的根节点指向v的根节点，合并两个集合

class NetworkEnvironment:
    def __init__(self, node_count=25, link_count=53):   
        self.graph = self._generate_geo_topology(node_count, link_count)
        self.source = "源节点"
        self.destination = "目标节点"
    
    def _draw_network(self,graph): # 复制的绘制网络图,要加self
        # 获取节点位置
        pos = nx.get_node_attributes(graph, 'pos')
    
        # 绘制图形
        plt.figure(figsize=(12, 8))
        nx.draw(graph, pos, with_labels=True, node_size=500, node_color='skyblue', font_size=10, font_weight='bold', edge_color='gray')
    
        # 添加标题
        plt.title("geo_top", fontsize=16) # 地理特征网络拓扑
        # plt.show() # show会阻塞当前线程，直到图形窗口关闭
    
    def _generate_geo_topology(self, n, m):
        G = nx.Graph()
        # 添加固定节点
        G.add_node("源节点", pos=(29.6469, 91.1175))
        G.add_node("目标节点", pos=(39.9042, 116.4074))
        # 生成其他节点坐标
        for i in range(n - 2):
            lat = random.uniform(28, 42)
            lon = random.uniform(90, 120)
            G.add_node(f"Node{i}", pos=(lat, lon))
        
        nodes = list(G.nodes())
        edges = []
        # 生成所有可能的边并计算距离
        for i in range(len(nodes)):
            for j in range(i + 1, len(nodes)):
                u = nodes[i]
                v = nodes[j]
                u_pos = G.nodes[u]['pos']
                v_pos = G.nodes[v]['pos']
                distance = self.haversine(u_pos[0], u_pos[1], v_pos[0], v_pos[1])
                edges.append((u, v, distance))
        
        # 按距离排序
        edges.sort(key=lambda x: x[2])   
        
        uf = UnionFind(nodes)
        selected_edges = []
        # 构建最小生成树确保连通
        for edge in edges: 
            u, v, dist = edge
            if uf.find(u) != uf.find(v): # 说明u和v不在同一连通分量
                uf.union(u, v)
                selected_edges.append((u, v))
                G.add_edge(u, v)         # 添加边到图中(我自己添加的).........
                if len(selected_edges) == n - 1:
                    break
        # self._draw_network(G)


        # 继续添加最近边直到达到m条
        selected_edges_set = set((u, v) for u, v in selected_edges) 
        # 遍历剩余边
        for edge in edges:
            if len(selected_edges) >= m:
                break
            u, v, _ = edge
            if (u, v) not in selected_edges_set:
                selected_edges.append((u, v))
                selected_edges_set.add((u, v))
        
        # 添加到图
        #G.add_edge("Tibet","Beijing")
        #selected_edges.append(("Tibet","Beijing"))  #需要把下一句的m改为m+1
        for u, v in selected_edges[:m]: 
            G.add_edge(u, v)

            # 设置地理相关的QoS参数
            base_delay = 10 if "Tibet" in (u, v) else 5  # 西藏链路基础延迟更高
            G[u][v].update({
                'delay': abs(np.random.normal(base_delay, 5)),  # 正态分布延迟
                'queue': np.random.randint(0, 15),
                'loss_rate': np.random.uniform(0.01, 0.2),
                'total_bw': 1000,  # Mbps
                'available_bw': np.random.uniform(100, 1000)
            })

        
        return G
    
    def haversine(self, lat1, lon1, lat2, lon2):
        R = 6371  # 地球半径(公里)
        dlat = math.radians(lat2 - lat1)
        dlon = math.radians(lon2 - lon1)
        a = (math.sin(dlat/2)**2 + 
             math.cos(math.radians(lat1)) *
             math.cos(math.radians(lat2)) * math.sin(dlon/2)**2)
        c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
        return R * c

    def get_neighbors(self, node):
        return list(self.graph.neighbors(node))
    
    def get_link_attrs(self, src, dst):  
        if not self.graph.has_edge(src, dst):   # 如果src和dst不相邻，报错
            raise ValueError(f"No edge between {src} and {dst}")   
        return self.graph[src][dst]

# 示例使用
#env = NetworkEnvironment(node_count=25, link_count=56)

# 绘制网络图
def draw_network(graph):
    # 获取节点位置
    pos = nx.get_node_attributes(graph, 'pos')
    
    # 绘制图形
    plt.figure(figsize=(12, 8))
    nx.draw(graph, pos, with_labels=True, node_size=500, node_color='skyblue', font_size=10, font_weight='bold', edge_color='gray')
    
    # 添加标题
    plt.title("geo_top_max", fontsize=16)#地理特征网络拓扑
    plt.draw() # plt.show()会阻塞当前线程，直到图形窗口关闭


# 调用绘制函数
#draw_network(env.graph)

#print(env.get_link_attrs("Tibet","Beijing")) #没有这个边会报错

# 所有节点圆形排列circular_layout
# pos = nx.circular_layout(env.graph)
# nx.draw(env.graph, pos, with_labels=True, node_color='skyblue', node_size=800)
# plt.show()
